{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.8 长短期记忆LSTM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import sys\n",
    "import d2lzh_pytorch as d2l\n",
    "\n",
    "#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "will use cpu\n"
     ]
    }
   ],
   "source": [
    "(corpus_indices, char_to_idx, idx_to_char, vocab_size) = \\\n",
    "    d2l.load_data_jay_lyrics()\n",
    "num_inputs, num_hiddens, num_outputs = vocab_size, 256, vocab_size\n",
    "print('will use', device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_params():\n",
    "    def _one(shape):\n",
    "        ts = torch.tensor(np.random.normal(0, 0.01, size=shape), \\\n",
    "                          device=device, dtype=torch.float32)\n",
    "        return torch.nn.Parameter(ts, requires_grad=True)\n",
    "    \n",
    "    def _three():\n",
    "        return (_one((num_inputs, num_hiddens)),\n",
    "               _one((num_hiddens, num_hiddens)),\n",
    "               torch.nn.Parameter(torch.zeros(num_hiddens, device=device,\n",
    "                                dtype=torch.float32), requires_grad=True))\n",
    "    \n",
    "    W_xi, W_hi, b_i = _three()    # 输入门参数\n",
    "    W_xf, W_hf, b_f = _three()    # 遗忘门参数\n",
    "    W_xo, W_ho, b_o = _three()    # 输出门参数\n",
    "    W_xc, W_hc, b_c = _three()    # 候选记忆细胞参数\n",
    "    \n",
    "    # 输出层参数\n",
    "    W_hq = _one((num_hiddens, num_outputs))\n",
    "    b_q = torch.nn.Parameter(torch.zeros(num_outputs,\n",
    "                device=device, dtype=torch.float32), requires_grad=True)\n",
    "    return nn.ParameterList([W_xi, W_hi, b_i, W_xf, W_hf, b_f, \n",
    "                            W_xo, W_ho, b_o, W_xc, W_hc, b_c, W_hq, b_q])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_lstm_state(batch_size, num_hiddens, device):\n",
    "    return (torch.zeros((batch_size, num_hiddens), device=device),\n",
    "           torch.zeros((batch_size, num_hiddens), device=device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lstm(inputs, state, params):\n",
    "    [W_xi, W_hi, b_i, W_xf, W_hf, b_f, W_xo, W_ho, b_o, W_xc, W_hc, b_c, W_hq, b_q] = params\n",
    "    (H, C) = state\n",
    "    outputs = []\n",
    "    for X in inputs:\n",
    "        I = torch.sigmoid(torch.matmul(X, W_xi) + torch.matmul(H, W_hi) + b_i)\n",
    "        F = torch.sigmoid(torch.matmul(X, W_xf) + torch.matmul(H, W_hf) + b_f)\n",
    "        O = torch.sigmoid(torch.matmul(X, W_xo) + torch.matmul(H, W_ho) + b_o)\n",
    "        C_tilda = torch.tanh(torch.matmul(X, W_xc) + torch.matmul(H, W_hc) + b_c)\n",
    "        C = F * C + I * C_tilda\n",
    "        H = O * C.tanh()\n",
    "        Y = torch.matmul(H, W_hq) + b_q\n",
    "        outputs.append(Y)\n",
    "    return outputs, (H, C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 40, perplexity 76.167146, time 11.97 sec\n",
      " - 分开 我想要你 我不要 我不要 我不不觉 我不不觉 我不不觉 我不不觉 我不不觉 我不不觉 我不不觉 我\n",
      " - 不分开 我想你你 我不要 我不要 我不不觉 我不不觉 我不不觉 我不不觉 我不不觉 我不不觉 我不不觉 我\n",
      "epoch 80, perplexity 7.343374, time 9.88 sec\n",
      " - 分开 我想你你的微笑每天都能看到  我知道这里很美但家乡的你更美   在着你的肩堡 你 在我胸口睡堡 像\n",
      " - 不分开 我想就你 我不要这熬  没有你你我有多多难多多恼   没有你烦堡 我 想和你的微笑 像说 你又很我\n",
      "epoch 120, perplexity 1.696393, time 11.03 sec\n",
      " - 分开 我想要你生微 天天看看运人 我都了我 你说 分数怎么停留 一直在停留 谁让它 说句的我 说说就通动\n",
      " - 不分开 为我不有你到出去 我想就没样 我想要恼  没后没你 我有开这样 我的天你 你手不离 我知再好生活 \n",
      "epoch 160, perplexity 1.635147, time 11.10 sec\n",
      " - 分开 我不能起的微笑 天通 却又常考倒着 说散 你想很久了吧? 我的认真败给黑色 别想 你想很久了吧? \n",
      " - 不分开 为我不觉到到痛一场悲有 你的画人 你种多 回  是却上的快快就像龙卷风 离家承受风圈来不及逃 我不\n"
     ]
    }
   ],
   "source": [
    "num_epochs, num_steps, batch_size, lr, clipping_theta = 160, 35, 16, 1e2, 1e-2\n",
    "pred_period, pred_len, prefixes = 40, 50, ['分开', '不分开']\n",
    "d2l.train_and_predict_rnn(lstm, get_params, init_lstm_state, num_hiddens,\n",
    "            vocab_size, device, corpus_indices, idx_to_char, char_to_idx,\n",
    "            False, num_epochs, num_steps, lr, clipping_theta, batch_size,\n",
    "            pred_period, pred_len, prefixes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
